package postgres

import (
	"context"
	"database/sql"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	log "ota/collector/logger"
	"ota/config"
	"ota/model"
)

var (
	dbOta               *gorm.DB
	dbAuth              *gorm.DB
	dbManger            *gorm.DB
	ContextTxKey        = struct{}{}
	ContextManagerTxKey = struct{}{}
	ContextAuthTxKey    = struct{}{}
)

func dbWithContext(ctx context.Context) *gorm.DB {
	return GetDB(ctx)
}

func GetManagerDB(ctx context.Context) *gorm.DB {
	tx, ok := ctx.Value(ContextManagerTxKey).(*gorm.DB)
	if ok {
		return tx
	}
	return dbManger.WithContext(ctx)
}

func GetAuthDB(ctx context.Context) *gorm.DB {
	tx, ok := ctx.Value(ContextAuthTxKey).(*gorm.DB)
	if ok {
		return tx
	}
	return dbAuth.WithContext(ctx)
}

func GetDB(ctx context.Context) *gorm.DB {
	tx, ok := ctx.Value(ContextTxKey).(*gorm.DB)
	if ok {
		return tx
	}
	return dbOta.WithContext(ctx)
}

func Init(connectAddr string, dbName string) func() {
	var (
		db    *gorm.DB
		err   error
		sqlDB *sql.DB
	)

	db, err = gorm.Open(postgres.Open(connectAddr), &gorm.Config{})
	if err != nil {
		log.Fatal(err.Error())
		return func() {}
	}

	if sqlDB, err = db.DB(); err != nil {
		log.Fatal(err.Error())
		return func() {}
	}

	sqlDB.SetMaxIdleConns(50)
	sqlDB.SetMaxOpenConns(500)
	if err = sqlDB.Ping(); err != nil {
		panic(err)
	}

	if config.Conf.Env == model.EnvDev {
		db = db.Debug()
	}
	for _, s := range migrates {
		s.Init()
	}

	switch dbName {
	case "ota":
		dbOta = db
		break
	case "auth":
		dbAuth = db
		break
	case "manager":
		dbManger = db
		break
	}

	return func() {
		if err = sqlDB.Close(); err != nil {
			log.Error(err.Error())
		}
	}
}

type IExecTx interface {
	ExecTx(ctx context.Context, fn func(ctx context.Context) error) error // 执行事务
}

type execTx struct {
}

func (a *execTx) ExecTx(ctx context.Context, fn func(ctx context.Context) error) error {
	return GetDB(ctx).Transaction(func(tx *gorm.DB) error {
		ctx = context.WithValue(ctx, ContextTxKey, tx)
		return fn(ctx)
	})

}
